## TFSA-2021-020: `CHECK`-fail in `tf.raw_ops.EncodePng`

### CVE Number
CVE-2021-29531

### Impact
An attacker can trigger a `CHECK` fail in PNG encoding by providing an empty
input tensor as the pixel data:

```python
import tensorflow as tf

image = tf.zeros([0, 0, 3])
image = tf.cast(image, dtype=tf.uint8)
tf.raw_ops.EncodePng(image=image)
```

This is because the
[implementation](https://github.com/tensorflow/tensorflow/blob/e312e0791ce486a80c9d23110841525c6f7c3289/tensorflow/core/kernels/image/encode_png_op.cc#L57-L60)
only validates that the total number of pixels in the image does not overflow.
Thus, an attacker can send an empty matrix for encoding. However, if the tensor
is empty, then the associated buffer is `nullptr`. Hence, when [calling
`png::WriteImageToBuffer`](https://github.com/tensorflow/tensorflow/blob/e312e0791ce486a80c9d23110841525c6f7c3289/tensorflow/core/kernels/image/encode_png_op.cc#L79-L93),
the first argument (i.e., `image.flat<T>().data()`) is `NULL`. This then
triggers the `CHECK_NOTNULL` in the [first line of
`png::WriteImageToBuffer`](https://github.com/tensorflow/tensorflow/blob/e312e0791ce486a80c9d23110841525c6f7c3289/tensorflow/core/lib/png/png_io.cc#L345-L349).

```cc
template <typename T>
bool WriteImageToBuffer(
    const void* image, int width, int height, int row_bytes, int num_channels,
    int channel_bits, int compression, T* png_string,
    const std::vector<std::pair<std::string, std::string> >* metadata) {
  CHECK_NOTNULL(image);
  ...
}
```

Since `image` is null, this results in `abort` being called after printing the
stacktrace. Effectively, this allows an attacker to mount a denial of service
attack.

### Patches
We have patched the issue in GitHub commit
[26eb323554ffccd173e8a79a8c05c15b685ae4d1](https://github.com/tensorflow/tensorflow/commit/26eb323554ffccd173e8a79a8c05c15b685ae4d1).

The fix will be included in TensorFlow 2.5.0. We will also cherrypick this
commit on TensorFlow 2.4.2, TensorFlow 2.3.3, TensorFlow 2.2.3 and TensorFlow
2.1.4, as these are also affected and still in supported range.

### For more information
Please consult [our security
guide](https://github.com/tensorflow/tensorflow/blob/master/SECURITY.md) for
more information regarding the security model and how to contact us with issues
and questions.

### Attribution
This vulnerability has been reported by Yakun Zhang and Ying Wang of Baidu
X-Team.
